// SPDX-FileCopyrightText: Copyright 2025-2025 深圳市同心圆网络有限公司
// SPDX-License-Identifier: GPL-3.0-only

package trace_cmd

import (
	"fmt"
	"sort"
	"time"

	"gitcode.com/opendragonfly/df_proto_gen_go.git/trace_api"
	"gitcode.com/opendragonfly/df_server/config"
	"gitcode.com/opendragonfly/df_server/utils"
	"github.com/spf13/cobra"
)

type SpanList []*trace_api.SpanInfo

func (l SpanList) Len() int {
	return len(l)
}

func (l SpanList) Swap(i, j int) {
	l[i], l[j] = l[j], l[i]
}

func (l SpanList) Less(i, j int) bool {
	return l[i].StartTimeStamp < l[j].StartTimeStamp
}

type SpanTreeNode struct {
	Span      *trace_api.SpanInfo
	ChildList []*SpanTreeNode
}

type SpanWithLevel struct {
	Span  *trace_api.SpanInfo
	Level uint32
}

type SpanWithLevelList struct {
	SpanList []*SpanWithLevel
}

func setupSpanTree(parentTreeNode *SpanTreeNode, spanList []*trace_api.SpanInfo) {
	for _, span := range spanList {
		if span.ParentSpanId == parentTreeNode.Span.SpanId {
			treeNode := &SpanTreeNode{
				Span:      span,
				ChildList: []*SpanTreeNode{},
			}
			parentTreeNode.ChildList = append(parentTreeNode.ChildList, treeNode)
			setupSpanTree(treeNode, spanList)
		}
	}
}

func walkSpanTree(parentTreeNode *SpanTreeNode, result *SpanWithLevelList, level uint32) {
	result.SpanList = append(result.SpanList, &SpanWithLevel{
		Span:  parentTreeNode.Span,
		Level: level,
	})
	for _, treeNode := range parentTreeNode.ChildList {
		walkSpanTree(treeNode, result, level+1)
	}
}

var listSpanCmd = &cobra.Command{
	Use: "listSpan",
	PreRunE: func(cmd *cobra.Command, args []string) error {
		if len(args) != 3 {
			return fmt.Errorf("need param [serviceName] [serviceVersion] [traceId]")
		}
		return nil
	},
	RunE: func(cmd *cobra.Command, args []string) error {
		cfg, err := config.ReadServerConfig()
		if err != nil {
			return err
		}
		if cfg.SuperToken == "" {
			return fmt.Errorf("miss super token in config")
		}
		conn, err := utils.ConnGrpcServer(fmt.Sprintf("127.0.0.1:%d", cfg.Port.ServicePort))
		if err != nil {
			return err
		}
		defer conn.Close()

		serviceName := args[0]
		serviceVersion := args[1]
		traceId := args[2]

		client := trace_api.NewTraceApiClient(conn)
		res, err := client.ListSpan(cmd.Context(), &trace_api.ListSpanRequest{
			AccessToken: cfg.SuperToken,
			Service: &trace_api.ServiceInfo{
				ServiceName:    serviceName,
				ServiceVersion: serviceVersion,
			},
			TraceId: traceId,
		})
		if err != nil {
			return err
		}
		if res.Code != trace_api.ListSpanResponse_CODE_OK {
			return fmt.Errorf("%s", res.ErrMsg)
		}
		//排序
		sort.Sort(SpanList(res.SpanList))
		//查找root span
		var rootTreeNode *SpanTreeNode
		for _, span := range res.SpanList {
			if span.ParentSpanId == "" {
				rootTreeNode = &SpanTreeNode{
					Span:      span,
					ChildList: []*SpanTreeNode{},
				}
				break
			}
		}
		if rootTreeNode == nil {
			return fmt.Errorf("miss root span")
		}
		//计算Span层级
		setupSpanTree(rootTreeNode, res.SpanList)
		result := &SpanWithLevelList{
			SpanList: []*SpanWithLevel{},
		}
		walkSpanTree(rootTreeNode, result, 0)
		//输出结果
		for _, spanWithLevel := range result.SpanList {
			for i := 0; i < int(spanWithLevel.Level); i++ {
				fmt.Print("  ")
			}
			fmt.Printf("%s startTime:%s,endTime:%s,consumeTime:%dms\n", spanWithLevel.Span.SpanName,
				time.UnixMilli(spanWithLevel.Span.StartTimeStamp).Format(time.DateTime),
				time.UnixMilli(spanWithLevel.Span.EndTimeStamp).Format(time.DateTime),
				spanWithLevel.Span.EndTimeStamp-spanWithLevel.Span.StartTimeStamp)
		}
		return nil
	},
}
